# start.py
import asyncio
import threading
import time
import sys
from mitmproxy.tools.dump import DumpMaster
from mitmproxy.options import Options
import signal
import os

# 确保正确导入capture_requests模块
try:
    import capture_requests
    print("✅ 成功导入 capture_requests 模块")
except ImportError as e:
    print(f"❌ 导入 capture_requests 模块失败: {e}")
    sys.exit(1)

class MitmproxyAPIStarter:
    def __init__(self, port=8080, ssl_port=8443, config_file="mitmproxy.conf"):
        self.config_file = config_file
        self.port = port
        self.ssl_port = ssl_port
        self.master = None
        self.loop = None
        self.thread = None
        self._load_config()

    def _load_config(self):
        """加载配置"""
        if not os.path.exists(self.config_file):
            print(f"⚠️ 配置文件 {self.config_file} 不存在，使用默认配置")
            return

        try:
            import configparser
            config = configparser.ConfigParser()
            config.read(self.config_file, encoding='utf-8')

            if 'settings' in config:
                settings = config['settings']
                if 'port' in settings:
                    self.port = int(settings['port'])
                if 'ssl_port' in settings:
                    self.ssl_port = int(settings['ssl_port'])

        except Exception as e:
            print(f"❌ 读取配置文件时出错: {e}")

    def _start_in_thread(self):
        """在线程中启动mitmproxy"""
        # 为这个线程创建新的事件循环并设置为当前线程的事件循环
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        try:
            # 配置选项，支持HTTPS
            opts = Options(
                listen_port=self.port,
                listen_host="127.0.0.1",
                ssl_insecure=True,  # 允许不安全的SSL连接
                certs=[],  # 可以指定自定义证书
            )

            # 如果配置了SSL端口，则同时监听HTTPS
            if self.ssl_port and self.ssl_port != self.port:
                opts.listen_port = self.port
                opts.listen_host = "127.0.0.1"

            # 尝试不同版本的DumpMaster构造函数
            try:
                # 新版本支持event_loop参数
                self.master = DumpMaster(opts, loop=self.loop)
                print("使用新版本DumpMaster构造函数")
            except TypeError:
                # 旧版本不支持event_loop参数
                self.master = DumpMaster(opts)
                print("使用旧版本DumpMaster构造函数")

            # 添加插件/脚本
            try:
                # 尝试使用addons方式
                if hasattr(capture_requests, 'addons') and capture_requests.addons:
                    for addon in capture_requests.addons:
                        self.master.addons.add(addon)
                    print("✅ 使用addons方式添加插件")
                else:
                    # 直接添加模块
                    self.master.addons.add(capture_requests)
                    print("✅ 使用模块方式添加插件")
            except AttributeError:
                # 如果没有addons属性，尝试其他方式
                print("⚠️ 无法通过标准方式添加插件")

            print(f"✅ mitmproxy已在端口 {self.port} 上启动")
            print(f"请将浏览器代理设置为: 127.0.0.1:{self.port}")
            print("💡 访问HTTPS网站时，需要安装并信任mitmproxy的CA证书")

            # 运行事件循环
            self.loop.run_until_complete(self.master.run())

        except Exception as e:
            print(f"❌ 启动mitmproxy时出错: {e}")
            import traceback
            traceback.print_exc()
        finally:
            if self.loop and not self.loop.is_closed():
                self.loop.close()

    def start(self):
        """启动mitmproxy"""
        self.thread = threading.Thread(target=self._start_in_thread, daemon=True)
        self.thread.start()

        # 等待启动
        time.sleep(2)

        # 检查是否成功启动
        if self.thread.is_alive():
            print("mitmproxy正在运行中...")
            return True
        else:
            print("mitmproxy启动失败")
            return False

    def shutdown(self):
        """关闭mitmproxy"""
        if self.master:
            # 在事件循环中执行关闭操作
            if self.loop and self.thread.is_alive():
                try:
                    future = asyncio.run_coroutine_threadsafe(self.master.shutdown(), self.loop)
                    future.result(timeout=5)
                except asyncio.TimeoutError:
                    print("关闭mitmproxy超时")
                except Exception as e:
                    print(f"关闭mitmproxy时出错: {e}")
            time.sleep(1)
        if self.thread:
            self.thread.join(timeout=5)

def main():
    # 获取配置文件路径
    config_file = "mitmproxy.conf"
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
        print(f"使用配置文件: {config_file}")

    proxy = MitmproxyAPIStarter(port=8080, ssl_port=8443)

    def signal_handler(sig, frame):
        print('\n正在停止mitmproxy...')
        proxy.shutdown()
        print("mitmproxy已停止")
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)

    if proxy.start():
        try:
            print("mitmproxy正在运行，按 Ctrl+C 停止")
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            signal_handler(signal.SIGINT, None)
    else:
        print("mitmproxy启动失败")

if __name__ == "__main__":
    main()
